from transformers.models.gpt2.modeling_gpt2 import GPT2Attention,GPT2FlashAttention2,GPT2Block,CausalLMOutputWithCrossAttentions
from transformers import GPT2Tokenizer, GPT2Model,GPT2LMHeadModel,GPT2Config
from typing import List, Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.distributed as dist
from ...pipeline.megatron import mpu
class GPT2Attention(GPT2Attention):

    def _merge_heads(self, x, hidden_size_per_partition):
        #[b,np,s,hn]->[b,s,np,hn]
        x = x.permute(0, 2, 1, 3).contiguous()
        new_shape = x.size()[:-2] + (hidden_size_per_partition,)
        return x.view(*new_shape)
    def _attn(self, query, key, value, attention_mask=None, head_mask=None):
        attn_weights = torch.matmul(query, key.transpose(-1, -2))#[b , np, sq, hn] * [b,np, hn,s] ->[b,np,s,s]
    
        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)
        attn_weights = attn_weights.type(value.dtype)
        attn_weights = self.attn_dropout(attn_weights)

        if head_mask is not None:
            attn_weights = attn_weights * head_mask
        #value [b , np, sq, hn]
        attn_output = torch.matmul(attn_weights, value) #[b,np,s,hn]
        return attn_output, attn_weights


    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
        if encoder_hidden_states is not None:
            if not hasattr(self, "q_attn"):
                raise ValueError(
                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
                )

            query = self.q_attn(hidden_states)
            key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
            attention_mask = encoder_attention_mask
        else:
            # query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
            mixed_x_layer = self.c_attn(hidden_states)  #[b,s,hidden_states*3]->[b,s,3*hn*np]
        tensor_parallel_size = 3 * self.embed_dim// mixed_x_layer.shape[-1]
        hidden_size_per_attention_head = self.embed_dim // self.num_heads   #hn 64
        hidden_size_per_partition = self.embed_dim //tensor_parallel_size   #hp 384
        num_attention_heads_per_partition = self.num_heads //tensor_parallel_size   #np,6
        new_tensor_shape = mixed_x_layer.size()[:-1] + (num_attention_heads_per_partition, 3 * hidden_size_per_attention_head)
        mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)   #[b, s, np, 3 * hn]
        query, key, value = mixed_x_layer.split(hidden_size_per_attention_head, dim=-1)  #3*[b, s, np, hn]

        query=query.permute(0, 2, 1, 3)#[b , np, sq, hn]
        key=key.permute(0, 2, 1, 3)#[b , np, sq, hn]
        value=value.permute(0, 2, 1, 3)#[b , np, sq, hn]
 
        if layer_past is not None:
            past_key, past_value = layer_past
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        # # 计算注意力得分

        attn_output,attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output, hidden_size_per_partition)

        attn_output = self.c_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)
        
        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs  # attn_output, present, (attentions)
class GPT2Block(GPT2Block):
    def __init__(self, config, layer_idx=None):
        super().__init__(config, layer_idx)
        attention_class = GPT2_ATTENTION_CLASSES[config._attn_implementation]
        self.attn = attention_class(config=config, layer_idx=layer_idx)
GPT2_ATTENTION_CLASSES = {
    "eager": GPT2Attention,
    "flash_attention_2": GPT2FlashAttention2,
}
class GPT2Model(GPT2Model):
    def __init__(self, config):
        super().__init__(config)
        self.h = nn.ModuleList([GPT2Block(config, layer_idx=i) for i in range(config.num_hidden_layers)])

class GPT2LMHeadModel(GPT2LMHeadModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.transformer = GPT2Model(config)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]

        # Set device for model parallelism
        if self.model_parallel:
            torch.cuda.set_device(self.transformer.first_device)
            hidden_states = hidden_states.to(self.lm_head.weight.device)

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = mpu.vocab_parallel_cross_entropy(shift_logits, shift_labels,self.config.vocab_size)
            loss = loss.mean()
            '''
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            '''
        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=loss,
            logits=lm_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
            cross_attentions=transformer_outputs.cross_attentions,
        )
